import numpy as np
from math_utils import Vector2


class RadiationSource:
    '''
    表示一个点状辐射源的类
    辐射强度会随着距离的增加而指数衰减
    '''

    def __init__(self, position: Vector2, intensity: float) -> None:
        '''
        初始化一个辐射源
        Args:
            position: 辐射源在二维空间中的坐标
            intensity: 辐射源的初始强度
        '''
        self.__position = position
        self.__intensity_4pi = intensity / (4 * np.pi)

    def get_dose(self, xx: np.ndarray, yy: np.ndarray) -> np.ndarray:
        '''计算由该辐射源产生的辐射剂量 (numpy 向量化计算)'''
        distance_squared = (xx - self.__position.x)**2 + (yy - self.__position.y)**2
        return self.__intensity_4pi / (distance_squared + 1e-7)


class RadiationField:
    '''表示一个由多个辐射源组成的辐射场'''

    def __init__(self, *sources: RadiationSource) -> None:
        '''初始化一个辐射场'''
        self.__sources = list(sources)

    def get_dose(self, xx: np.ndarray, yy: np.ndarray) -> np.ndarray:
        '''计算由整个辐射场产生的总辐射剂量, 取最小的99%, 归一化 (numpy 向量化计算)'''
        # 初始化剂量数组为零
        dose = np.zeros_like(xx)

        # 累加所有辐射源的剂量
        for source in self.__sources:
            dose += source.get_dose(xx, yy)

        # 取98%分位数作为最大值
        set_max = np.percentile(dose, q=98)
        dose[dose > set_max] = set_max

        # 归一化处理
        dose -= dose.min()
        dose /= dose.max()
        return dose

    def append(self, *sources: RadiationSource) -> None:
        '''向辐射场中添加一个或多个新的辐射源'''
        self.__sources.extend(sources)

    def __getitem__(self, index: int) -> RadiationSource:
        '''通过索引获取辐射场中的一个辐射源'''
        return self.__sources[index]

    def __setitem__(self, index: int, value: RadiationSource) -> None:
        '''通过索引设置或替换辐射场中的一个辐射源'''
        self.__sources[index] = value

    def __len__(self) -> int:
        '''获取辐射场中辐射源的数量'''
        return len(self.__sources)


if __name__ == '__main__':

    # 创建辐射源
    source1 = RadiationSource(Vector2(0, 0), intensity=100)
    source2 = RadiationSource(Vector2(5, 0), intensity=80)
    source3 = RadiationSource(Vector2(0, 5), intensity=60)

    # 创建辐射场
    field = RadiationField(source1, source2, source3)
